import gc
import os
import pickle
import random

import numpy as np
import soundfile
from PIL import Image
from scipy.signal import resample
from ssqueezepy import ssq_stft

from Artear.utils import util


def audio2params_data_gen(data_source, config, batch_size, input_type='audio', output_type='params'):
    
    from utils.simpledaw import SimpleDAW
    presets = util.get_files(data_source['preset'], 'fxp')
    midis = util.get_files(data_source['midi'], 'mid')
    daw = SimpleDAW(plugin=config['plugin_path'],
                    sample_rate=config['source_sample_rate'], bpm=config['bpm'])

    while True:
        spec_data, param_data, midi_data, raw_audio_data = [], [], [], []

        for i in range(batch_size):
            params = np.array(daw.get_params(random.choice(presets)))
            param_data.append(params.reshape((1,) + params.shape))
            midi = random.choice(midis)
            daw.load_midi(midi)
            midi_data.append(midi)
            raw = daw.render(config['duration'])
            raw_audio_data.append(raw.reshape((1,) + raw.shape))
            dat_0 = resample(raw[:, 0], int(
                raw.shape[-2] / config['compress_rate']))
            dat_1 = resample(raw[:, 1], int(
                raw.shape[-2] / config['compress_rate']))
            # spec_0 = librosa.feature.melspectrogram(raw[:, 0], config['source_sample_rate'], hop_length=512)
            # spec_1 = librosa.feature.melspectrogram(raw[:, 1], config['source_sample_rate'], hop_length=512)
            spec_0, *_ = ssq_stft(dat_0)
            spec_1, *_ = ssq_stft(dat_1)

            spec_0 = np.abs(spec_0)
            spec_1 = np.abs(spec_1)

            '''img_0 = Image.fromarray(spec_0)
            img_0 = img_0.resize((224, 224))
            spec_0 = np.array(img_0)
            img_1 = Image.fromarray(spec_1)
            img_1 = img_0.resize((224, 224))
            spec_1 = np.array(img_1)'''

            def resize(spec_):
                if spec_.shape[0] % 2 != 0:
                    img_ = Image.fromarray(spec_)
                    img_ = img_.resize((spec_.shape[0] - 1, spec_.shape[1]))
                    spec_ = np.array(img_)
                if spec_.shape[1] % 2 != 0:
                    img_ = Image.fromarray(spec_)
                    img_ = img_.resize((spec_.shape[0], spec_.shape[1] - 1))
                    spec_ = np.array(img_)
                return spec_

            spec_0 = resize(spec_0)
            spec_1 = resize(spec_1)
            '''spec_0 = spec_0.transpose((1, 0))
            spec_1 = spec_1.transpose((1, 0))'''
            spec = np.stack([spec_0, spec_1], axis=-1)
            # print(spec.shape)
            spec_data.append(spec.reshape((1,) + spec.shape))

        # print(raw_audio_data[0].shape, spec_data[0].shape, param_data[0].shape)
        if input_type == 'audio':
            in_data = spec_data
        elif input_type == 'params':
            in_data = param_data
        else:
            in_data = spec_data

        if output_type == 'audio':
            out_data = spec_data
        elif output_type == 'params':
            out_data = param_data
        else:
            out_data = param_data

        config['src_len'] = raw_audio_data[0].shape[1]
        config['spec_len'] = spec_data[0].shape[1]

        yield {
            'spec': spec_data, 'params': param_data, 'midi': midi_data, 'raw_audio': raw_audio_data,
            'in': in_data, 'out': out_data,
        }

        del spec_data, param_data, midi_data, raw_audio_data
        gc.collect()


def get_raw_audio_data_to_file(data_source, config, batch_size, size, input_type='audio', output_type='params'):

    from utils.simpledaw import SimpleDAW
    presets = util.get_files(data_source['preset'], 'fxp')
    midis = util.get_files(data_source['midi'], 'mid')
    daw = SimpleDAW(plugin=config['plugin_path'],
                    sample_rate=config['source_sample_rate'], bpm=config['bpm'])
    
    dataset = []
    for i in range(size):
        spec_data, param_data, midi_data, raw_audio_data = [], [], [], []
    
        for j, preset in enumerate(presets[:batch_size]):
            print(j)
        
            params = np.array(daw.get_params(preset))
            param_data.append(params.reshape((1,) + params.shape))
            midi = random.choice(midis)
            daw.load_midi(midi)
            midi_data.append(midi)
            raw = daw.render(config['duration'])
            with open(config['save_dir'] + '\\' + ''.join(os.path.basename(preset).split('.')[:-1]) + '.wav', 'wb') as file:
                soundfile.write(file, raw, int(config['source_sample_rate']))
            raw_audio_data.append(raw.reshape((1,) + raw.shape))

        dataset.append({
            'spec': spec_data, 'params': param_data, 'midi': midi_data, 'raw_audio': raw_audio_data,
        })
        
    with open(data_source['pickled_data'], 'wb') as file:
        pickle.dump(dataset, file)
    


def audio2params_data_gen_from_file(data_source, config, batch_size, input_type='audio', output_type='params'):
    spec_data, param_data, midi_data, raw_audio_data = [], [], [], []
    with open(data_source['pickled_data'], 'wb') as file:
        dataset = pickle.load(file)

    while True:
        for data in dataset:

            spec_data, param_data, midi_data, raw_audio_data = [], [], [], data['raw_audio']
        
            for raw in raw_audio_data:
                dat_0 = resample(raw[:, 0], int(
                    raw.shape[-2] / config['compress_rate']))
                dat_1 = resample(raw[:, 1], int(
                    raw.shape[-2] / config['compress_rate']))
                # spec_0 = librosa.feature.melspectrogram(raw[:, 0], config['source_sample_rate'], hop_length=512)
                # spec_1 = librosa.feature.melspectrogram(raw[:, 1], config['source_sample_rate'], hop_length=512)
                spec_0, *_ = ssq_stft(dat_0)
                spec_1, *_ = ssq_stft(dat_1)

                spec_0 = np.abs(spec_0)
                spec_1 = np.abs(spec_1)

                def resize(spec_):
                    if spec_.shape[0] % 2 != 0:
                        img_ = Image.fromarray(spec_)
                        img_ = img_.resize((spec_.shape[0] - 1, spec_.shape[1]))
                        spec_ = np.array(img_)
                    if spec_.shape[1] % 2 != 0:
                        img_ = Image.fromarray(spec_)
                        img_ = img_.resize((spec_.shape[0], spec_.shape[1] - 1))
                        spec_ = np.array(img_)
                    return spec_

                spec_0 = resize(spec_0)
                spec_1 = resize(spec_1)
                spec = np.stack([spec_0, spec_1], axis=-1)
                # print(spec.shape)
                spec_data.append(spec.reshape((1,) + spec.shape))

            # print(raw_audio_data[0].shape, spec_data[0].shape, param_data[0].shape)
            if input_type == 'audio':
                in_data = spec_data
            elif input_type == 'params':
                in_data = param_data
            else:
                in_data = spec_data

            if output_type == 'audio':
                out_data = spec_data
            elif output_type == 'params':
                out_data = param_data
            else:
                out_data = param_data

            config['src_len'] = raw_audio_data[0].shape[1]
            config['spec_len'] = spec_data[0].shape[1]

            yield {
                'spec': spec_data, 'params': param_data, 'midi': midi_data, 'raw_audio': raw_audio_data,
                'in': in_data, 'out': out_data,
            }

            del spec_data, param_data, midi_data, raw_audio_data
            gc.collect()
        

def data_gen_from_file(data_source):
    with open(data_source['pickled_data'], 'wb') as file:
        dataset = pickle.load(file)
    
    while True:
        for data in dataset:
            yield data
        
        
def get_params2params_data_to_file(data_source, config, batch_size, size):
    
    from utils.simpledaw import SimpleDAW
    presets = util.get_files(data_source['preset'], 'fxp')
    daw = SimpleDAW(plugin=config['plugin_path'],
                    sample_rate=config['source_sample_rate'], bpm=config['bpm'])
    
    dataset = []

    for i in range(size):
        param_data = []
        for j in range(batch_size):
            params = np.array(daw.get_params(random.choice(presets)))
            param_data.append(params.reshape((1,) + params.shape))
         
        dataset.append({
            'spec': None, 'params': param_data, 'midi': None, 'raw_audio': None,
            'in': param_data, 'out': None,
        })
    
    with open(data_source['pickled_data'], 'wb') as file:
        pickle.dump(dataset, file)
    


if __name__ == '__main__':
    data_source = {
        'preset': r"B:\muse_repo\Presets\_Legacy",
        'midi': r"B:\muse_repo\MIDI",
        # 'pickled_data': r"/home/dajun/data.pickled",
        'pickled_data': r"C:/dev_spa/DMuse/data.pickled",
    }

    config = {
        'model_type': 'audio2params',
        'plugin_path': 'C:/VST/64bit/Sylenth1.dll',
        'save_dir': r'C:\dev_spa\DMuse',
        'duration': 4.,
        'source_sample_rate': 4410.,
        'compress_rate': 200.,
        'bpm': 120.,
    }
    batch_size = 4
    size = 2000
    

    # data = get_raw_audio_data_to_file(data_source, config, size)
    
    # data_gen = audio2params_data_gen(data_sources, config, size)
    # data = next(data_gen)
    
    get_params2params_data_to_file(data_source, config, batch_size, size)
    
